package com.limit.aspect.core;

import com.limit.annotation.LeakyBucketLimit;
import com.limit.annotation.SlidingWindowLimit;
import com.limit.annotation.TokenBucketLimit;
import com.limit.constant.RedisKeyConstant;
import com.limit.constant.RedisLuaConstant;
import com.limit.enums.RateLimitStrategyEnum;
import com.limit.exception.RequestLimitException;
import com.limit.util.ServletUtil;
import com.limit.util.UserUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.StringJoiner;
import java.util.function.BiConsumer;

/**
 * @title: RateLimitingStrategy
 * @Author xu-xiang
 * @Date: 2024/3/25 16:02
 * @description: 限流策略
 */
@Service
public class RateLimitingStrategyService {

    private static final Logger logger = LoggerFactory.getLogger(RateLimitingStrategyService.class);

    // 此处用StringRedisTemplate, 以保证lua脚本可以正确识别ARGV
    // 除非RedisTemplate重写了序列化与反序列化机制, 否则lua参数可能为nil
    @Resource(name = "stringRedisTemplate")
    private StringRedisTemplate redisTemplate;

    private static Map<Class<? extends Annotation>, BiConsumer<Annotation, Method>> STRATEGY_MAP;

    private static final DefaultRedisScript<Long> TOKEN_BUCKET_LIMIT_LUA_SCRIPT = new DefaultRedisScript<>();
    private static final DefaultRedisScript<Long> LEAKY_BUCKET_LIMIT_LUA_SCRIPT = new DefaultRedisScript<>();
    private static final DefaultRedisScript<Long> SLIDING_WINDOW_LIMIT_LUA_SCRIPT = new DefaultRedisScript<>();

    static {
        // 初始化lua脚本Source
        TOKEN_BUCKET_LIMIT_LUA_SCRIPT.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/tokenBucketLimit.lua")));
        LEAKY_BUCKET_LIMIT_LUA_SCRIPT.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/leakyBucketLimit.lua")));
        SLIDING_WINDOW_LIMIT_LUA_SCRIPT.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/slidingWindowLimit.lua")));
        TOKEN_BUCKET_LIMIT_LUA_SCRIPT.setResultType(Long.class);
        LEAKY_BUCKET_LIMIT_LUA_SCRIPT.setResultType(Long.class);
        SLIDING_WINDOW_LIMIT_LUA_SCRIPT.setResultType(Long.class);
    }

    @PostConstruct
    public void register() {
        Map<Class<? extends Annotation>, BiConsumer<Annotation, Method>> annotationMap = new HashMap<>(4);

        // 令牌桶限流策略
        BiConsumer<Annotation, Method> tokenBucketLimitConsumer = (annotation, method) -> {
            TokenBucketLimit tokenBucketLimit = (TokenBucketLimit) annotation;
            this.annotationCheck(annotation);

            RateLimitStrategyEnum[] strategy = tokenBucketLimit.strategy();
            String redisKey = getRedisKey(method, strategy);

            long curMillis = System.currentTimeMillis();
            Long result = redisTemplate.execute(TOKEN_BUCKET_LIMIT_LUA_SCRIPT,
                    Collections.singletonList(redisKey),  // 限流key
                    String.valueOf(tokenBucketLimit.consumptionPerRequest()),  //  每次消耗令牌数量
                    String.valueOf(curMillis), // 当前时间戳毫秒值
                    String.valueOf(tokenBucketLimit.bucketCapacity()),  // 最大令牌桶数
                    String.valueOf(tokenBucketLimit.period())  // 令牌有效期
            );

            Assert.notNull(result, "lua script execute error");
            if (result != RedisLuaConstant.PASS) {
                logger.error("请求过于频繁,当前允许最大请求频率:{}次/秒", tokenBucketLimit.bucketCapacity() / tokenBucketLimit.period());
                throw new RequestLimitException();
            }
        };

        // 漏桶限流策略
        BiConsumer<Annotation, Method> leakyBucketLimitConsumer = (annotation, method) -> {
            LeakyBucketLimit leakyBucketLimit = (LeakyBucketLimit) annotation;
            this.annotationCheck(annotation);

            RateLimitStrategyEnum[] strategy = leakyBucketLimit.strategy();
            String redisKey = getRedisKey(method, strategy);

            long curMillis = System.currentTimeMillis();
            Long result = redisTemplate.execute(LEAKY_BUCKET_LIMIT_LUA_SCRIPT,
                    Collections.singletonList(redisKey),  // 限流key
                    String.valueOf(leakyBucketLimit.leakyCapacity()),  // 最大桶容量
                    String.valueOf(leakyBucketLimit.dripRate()),  //  漏桶流出速率
                    String.valueOf(curMillis)  // 当前时间戳毫秒值
            );

            Assert.notNull(result, "lua script execute error");
            if (result != RedisLuaConstant.PASS) {
                logger.error("请求过于频繁,请稍候再试");
                throw new RequestLimitException();
            }
        };

        // 滑动窗口限流策略
        BiConsumer<Annotation, Method> slidingWindowLimitConsumer = (annotation, method) -> {
            SlidingWindowLimit slidingWindowLimit = (SlidingWindowLimit) annotation;
            this.annotationCheck(annotation);

            RateLimitStrategyEnum[] strategy = slidingWindowLimit.strategy();
            String redisKey = getRedisKey(method, strategy);

            long curMillis = System.currentTimeMillis();
            Long result = redisTemplate.execute(SLIDING_WINDOW_LIMIT_LUA_SCRIPT,
                    Collections.singletonList(redisKey),  // 限流key
                    String.valueOf(slidingWindowLimit.maxTimeInterval()),  // 在X秒内
                    String.valueOf(slidingWindowLimit.maxRequestCount()),  // 最多请求次数
                    String.valueOf(curMillis)  // 当前时间戳毫秒值
            );

            Assert.notNull(result, "lua script execute error");
            if (result != RedisLuaConstant.PASS) {
                logger.error("请求过于频繁,当前允许最大请求频率:{}次/{}秒", slidingWindowLimit.maxRequestCount(), slidingWindowLimit.maxTimeInterval());
                throw new RequestLimitException();
            }
        };

        annotationMap.put(TokenBucketLimit.class, tokenBucketLimitConsumer);
        annotationMap.put(LeakyBucketLimit.class, leakyBucketLimitConsumer);
        annotationMap.put(SlidingWindowLimit.class, slidingWindowLimitConsumer);
        STRATEGY_MAP = Collections.unmodifiableMap(annotationMap);
    }

    public void doRedisLua(Annotation limitAnnotation, Method method) {
        BiConsumer<Annotation, Method> annotationConsumer = STRATEGY_MAP.get(limitAnnotation.annotationType());
        if (annotationConsumer != null) {
            annotationConsumer.accept(limitAnnotation, method);
        } else {
            throw new IllegalArgumentException("No strategy found for annotation: " + limitAnnotation);
        }
    }

    /**
     * 根据策略构建redisKey
     *
     * @param method        请求方法
     * @param strategyEnums 策略枚举
     * @return redisKey
     */
    private String getRedisKey(Method method, RateLimitStrategyEnum[] strategyEnums) {
        StringJoiner redisKeyJoiner = new StringJoiner(":");
        redisKeyJoiner.add(RedisKeyConstant.RATE_LIMIT);

        for (RateLimitStrategyEnum strategyEnum : strategyEnums) {
            if (strategyEnum == RateLimitStrategyEnum.IP) {
                redisKeyJoiner.add(ServletUtil.getClientIpAddr());
            } else if (strategyEnum == RateLimitStrategyEnum.METHOD) {
                redisKeyJoiner.add(method.getName());
            } else if (strategyEnum == RateLimitStrategyEnum.USER) {
                redisKeyJoiner.add(UserUtil.getCurrentUserId());
            }
        }

        return redisKeyJoiner.toString();
    }

    /**
     * 注解参数校验
     */
    private void annotationCheck(Annotation annotation) {
        if (annotation == null) {
            throw new NullPointerException();
        }

        if (annotation.annotationType() == TokenBucketLimit.class) {
            TokenBucketLimit tokenBucketLimit = (TokenBucketLimit) annotation;
            Assert.notNull(tokenBucketLimit, "无法获取tokenBucketLimit");
            Assert.isTrue(tokenBucketLimit.bucketCapacity() > 0, "bucketCapacity 必须大于0");
            Assert.isTrue(tokenBucketLimit.period() > 0, "period 必须大于0");
        } else if (annotation.annotationType() == LeakyBucketLimit.class) {
            LeakyBucketLimit leakyBucketLimit = (LeakyBucketLimit) annotation;
            Assert.notNull(leakyBucketLimit, "无法获取tokenBucketLimit");
            Assert.isTrue(leakyBucketLimit.leakyCapacity() > 0, "leakyCapacity 必须大于0");
            Assert.isTrue(leakyBucketLimit.dripRate() > 0, "dripRate 必须大于0");
        } else if (annotation.annotationType() == SlidingWindowLimit.class) {
            SlidingWindowLimit slidingWindowLimit = (SlidingWindowLimit) annotation;
            Assert.notNull(slidingWindowLimit, "无法获取slidingWindowLimit");
            Assert.isTrue(slidingWindowLimit.maxTimeInterval() > 0, "maxTimeInterval 必须大于0");
            Assert.isTrue(slidingWindowLimit.maxRequestCount() > 0, "maxRequestCount 必须大于0");
        }
    }
}
